package com.leon.datalink.core.storage.impl;


import com.leon.datalink.core.evn.EnvUtil;
import com.leon.datalink.core.storage.DataStorage;
import com.leon.datalink.core.utils.DiskUtils;
import com.leon.datalink.core.utils.Loggers;
import com.leon.datalink.core.utils.ProtostuffUtils;

import java.io.File;
import java.io.IOException;
import java.sql.*;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.ConcurrentHashMap;
import java.util.Map;

/**
 * @ClassName ObjectStorage
 * @Description H2 Database Storage (替代LMDB)
 * @Author Leon
 * @Date 2023年7月31日16:52:52
 * @Version V1.0
 **/
public class ObjectStorage<T> implements DataStorage<T> {

    private static Connection connection;
    private static final Map<String, Boolean> tableCreated = new ConcurrentHashMap<>();

    private final Class<T> clazz;
    private final String tableName;

    public ObjectStorage(Class<T> clazz) throws IOException {
        synchronized (ObjectStorage.class) {
            if (connection == null) {
                try {
                    File path = new File(EnvUtil.getStoragePath());
                    DiskUtils.forceMkdir(path);
                    String dbPath = "jdbc:h2:" + path.getAbsolutePath() + "/datalink;AUTO_SERVER=TRUE;DB_CLOSE_DELAY=-1";
                    connection = DriverManager.getConnection(dbPath, "sa", "");
                    connection.setAutoCommit(true);
                    Loggers.CORE.info("Create H2 database connection: {}", dbPath);
                } catch (SQLException e) {
                    throw new IOException("Failed to create H2 connection", e);
                }
            }
        }
        this.clazz = clazz;
        this.tableName = clazz.getSimpleName().toLowerCase();
        createTableIfNotExists();
    }

    private void createTableIfNotExists() {
        if (tableCreated.containsKey(tableName)) {
            return;
        }

        String createTableSQL = String.format(
            "CREATE TABLE IF NOT EXISTS \"%s\" (" +
            "\"key\" VARCHAR(255) PRIMARY KEY, " +
            "\"value\" BLOB, " +
            "\"created_at\" TIMESTAMP DEFAULT CURRENT_TIMESTAMP, " +
            "\"updated_at\" TIMESTAMP DEFAULT CURRENT_TIMESTAMP" +
            ")", tableName
        );

        try (Statement stmt = connection.createStatement()) {
            stmt.execute(createTableSQL);
            tableCreated.put(tableName, true);
            Loggers.CORE.info("Created H2 table: {}", tableName);
        } catch (SQLException e) {
            Loggers.CORE.error("Failed to create table: " + tableName, e);
        }
    }

    @Override
    public T get(String key) {
        String sql = String.format("SELECT \"value\" FROM \"%s\" WHERE \"key\" = ?", tableName);

        try (PreparedStatement pstmt = connection.prepareStatement(sql)) {
            pstmt.setString(1, key);

            try (ResultSet rs = pstmt.executeQuery()) {
                if (rs.next()) {
                    byte[] data = rs.getBytes("value");
                    return ProtostuffUtils.deserialize(data, clazz);
                }
            }
        } catch (SQLException e) {
            Loggers.CORE.error("Failed to get key: " + key, e);
        }
        return null;
    }

    @Override
    public void put(String key, T value) {
        byte[] data = ProtostuffUtils.serialize(value);

        String sql = String.format(
            "MERGE INTO \"%s\" (\"key\", \"value\", \"updated_at\") VALUES (?, ?, CURRENT_TIMESTAMP)",
            tableName
        );

        try (PreparedStatement pstmt = connection.prepareStatement(sql)) {
            pstmt.setString(1, key);
            pstmt.setBytes(2, data);
            pstmt.executeUpdate();
        } catch (SQLException e) {
            Loggers.CORE.error("Failed to put key: " + key, e);
        }
    }

    @Override
    public void delete(String key) {
        String sql = String.format("DELETE FROM \"%s\" WHERE \"key\" = ?", tableName);

        try (PreparedStatement pstmt = connection.prepareStatement(sql)) {
            pstmt.setString(1, key);
            pstmt.executeUpdate();
        } catch (SQLException e) {
            Loggers.CORE.error("Failed to delete key: " + key, e);
        }
    }

    @Override
    public List<String> getKeys() {
        List<String> result = new ArrayList<>();
        String sql = String.format("SELECT \"key\" FROM \"%s\"", tableName);

        try (Statement stmt = connection.createStatement();
             ResultSet rs = stmt.executeQuery(sql)) {

            while (rs.next()) {
                result.add(rs.getString("key"));
            }
        } catch (SQLException e) {
            Loggers.CORE.error("Failed to get keys", e);
        }
        return result;
    }

    @Override
    public List<T> getValues() {
        List<T> result = new ArrayList<>();
        String sql = String.format("SELECT \"value\" FROM \"%s\"", tableName);

        try (Statement stmt = connection.createStatement();
             ResultSet rs = stmt.executeQuery(sql)) {

            while (rs.next()) {
                byte[] data = rs.getBytes("value");
                result.add(ProtostuffUtils.deserialize(data, clazz));
            }
        } catch (SQLException e) {
            Loggers.CORE.error("Failed to get values", e);
        }
        return result;
    }

    @Override
    public long count() {
        String sql = String.format("SELECT COUNT(*) FROM \"%s\"", tableName);

        try (Statement stmt = connection.createStatement();
             ResultSet rs = stmt.executeQuery(sql)) {

            if (rs.next()) {
                return rs.getLong(1);
            }
        } catch (SQLException e) {
            Loggers.CORE.error("Failed to count records", e);
        }
        return 0;
    }

}
